import sensor
from machine import UART
import image
import lcd
import KPU as kpu
import time
from Maix import FPIOA, GPIO
import gc
from fpioa_manager import fm
from board import board_info
import utime

color_R = (255, 0, 0)
color_G = (0, 255, 0)
color_B = (0, 0, 255)
class_IDs = ['no_mask', 'mask']
############### config #################
class_num = 3
sample_num = (class_num*5)
THRESHOLD = 11
class_names = ['class1', 'class2', 'class3']
board_cube = 0
########################################
fm.register(21, fm.fpioa.GPIOHS0)   #indepent
fm.register(22, fm.fpioa.GPIOHS1)    #mask
fm.register(32, fm.fpioa.GPIOHS2)    #indepent   ������Ƭ
fm.register(23, fm.fpioa.GPIOHS3)     #����ʶ����������ȡface learn
fm.register(5, fm.fpioa.UART1_TX, force=True)
fm.register(4, fm.fpioa.UART1_RX, force=True)
uart_A = UART(UART.UART1, 115200, 8, 1, 0, timeout=1000, read_buf_len=4096)
KEY0=GPIO(GPIO.GPIOHS0, GPIO.IN, GPIO.PULL_UP)
KEY1=GPIO(GPIO.GPIOHS1, GPIO.IN, GPIO.PULL_UP)
key=GPIO(GPIO.GPIOHS2, GPIO.IN, GPIO.PULL_UP)
key_gpio=GPIO(GPIO.GPIOHS3, GPIO.IN, GPIO.PULL_UP)
start_processing = False #������ȫ�ֱ�����ԭ����ʲô��
BOUNCE_PROTECTION = 200
def set_key_state(*_):   #ΪʲôҪ����������һ���������а����ж�
    global start_processing
    start_processing = True
    utime.sleep_ms(BOUNCE_PROTECTION)  #�����жϵ������ǽ���������Ƭ��������������ȡ
key_gpio.irq(set_key_state, GPIO.IRQ_FALLING,GPIO.WAKEUP_NOT_SUPPORT)
###########
def drawConfidenceText(image, rol, classid, value):
    text = ""
    _confidence = int(value * 100)

    if classid == 1:
        uart_A.write('1\r\n')
        text = 'mask: ' + str(_confidence) + '%'
        color_text=color_G
    else:
        uart_A.write('2\r\n')
        text = 'no_mask: ' + str(_confidence) + '%'
        color_text=color_R
    image.draw_string(rol[0], rol[1], text, color=color_text, scale=2.5)
    #############
def mask_learn():
    lcd.init()
    sensor.reset(dual_buff=True)   #��������ͷ��˫����ģʽ
    sensor.set_pixformat(sensor.RGB565)
    sensor.set_framesize(sensor.QVGA)
    sensor.set_hmirror(1)    #����ˮƽ��ת
    sensor.run(1)    #��ʼץȡͼ��

    task = kpu.load("/sd/mask.smodel")   #����sd��������ģ��

    anchor = (0.1606, 0.3562, 0.4712, 0.9568, 0.9877, 1.9108, 1.8761, 3.5310, 3.4423, 5.6823)
    _ = kpu.init_yolo2(task, 0.5, 0.3, 5, anchor)
    img_lcd = image.Image()  #����һ���հ�ͼ��//��������trueֱ�Ӽ��ؽ�֡�������� FALSE ���ؽ�python�ڵĶ�����

    clock = time.clock()  #����һ��ʱ�Ӷ���
    while (True):
        clock.tick()  #��ʼ��������ʱ��
        img = sensor.snapshot()    #��������ͷ����һ����Ƭ
        code = kpu.run_yolo2(task, img)
        if code:
            totalRes = len(code)

            for item in code:
                confidence = float(item.value()) #�����ĵ���Ƭ������Ϣ��ģ�ͱȶԣ�����������
                itemROL = item.rect()  #rect�Ǿ��Σ���������
                classID = int(item.classid())

                if confidence < 0.52:
                    _ = img.draw_rectangle(itemROL, color=color_B, tickness=5) #tick���ߴ�
                    continue

                if classID == 1 and confidence > 0.65:
                    _ = img.draw_rectangle(itemROL, color_G, tickness=5)
                    if totalRes == 1:
                        drawConfidenceText(img, (0, 0), 1, confidence)
                else:
                    _ = img.draw_rectangle(itemROL, color=color_R, tickness=5)
                    if totalRes == 1:

                     drawConfidenceText(img, (0, 0), 0, confidence)

        _ = lcd.display(img)
    _ = kpu.deinit(task)

###########
def draw_string(img, x, y, text, color, scale, bg=None ):
    if bg:
        img.draw_rectangle(x-2,y-2, len(text)*8*scale+4 , 16*scale, fill=True, color=bg)
    img = img.draw_string(x, y, text, color=color,scale=scale)
    return img

def indepent_learn():
    sensor.reset()
    sensor.set_pixformat(sensor.RGB565)
    sensor.set_framesize(sensor.QVGA)
    sensor.set_windowing((224, 224))
    if board_cube == 1:
        sensor.set_vflip(True)
        sensor.set_hmirror(True)
        lcd.init(type=2)  #lcd����
        lcd.rotation(2)   #����lcd����
    else:
        lcd.init()
    try:
        del model
    except Exception:
        pass
    try:
        del classifier
    except Exception:
        pass
    gc.collect()   #������������
    model = kpu.load("/sd/learn.smodel")
    classifier = kpu.classifier(model, class_num, sample_num)

    cap_num = 0
    train_status = 0  #ѵ����־λΪ0
    last_cap_time = 0
    last_btn_status = 1
    while 1:
        img = sensor.snapshot()  #����һ����Ƭ,��Ϊ����
        if board_cube:
            img = img.rotation_corr(z_rotation=90)  #��תͼ��
            img.pix_to_ai() #��Ϊʹ����img����������
            #������Ҫʹ�������ӿ�kpu�����ٶȣ���Ϊͼ����������kpu�ϵ�
        # capture img  ����ͼ��
        if train_status == 0:
            if key.value() == 0:
                time.sleep_ms(20)   #��������
                if key.value() == 0 and (last_btn_status == 1) and (time.ticks_ms() - last_cap_time > 500):  #�������������Ż�ִ�����г���������һ���������ⲻ��
                    last_btn_status = 0   #��0��Ŀ����
                    last_cap_time = time.ticks_ms() #ΪʲôҪҪ��Ŀǰ��ʱ����������һ����
                    if cap_num < class_num:
                        index = classifier.add_class_img(img)
                        cap_num += 1#ע�Ⱑ����ִ����֮���Ͳ�����ִ��elif�ж������Ķ�����
                        print("add class img:", index)
                    elif cap_num < class_num + sample_num:
                        index = classifier.add_sample_img(img)
                        cap_num += 1  #������Ҫ��Ϊ�˲ɼ����ݼ�����ѵ��
                        print("add sample img:", index)
                else:
                    img = draw_string(img, 2, 200, "release boot key please", color=lcd.WHITE,scale=1, bg=lcd.RED)
            else:
                time.sleep_ms(30)
                if key.value() == 1 and (last_btn_status == 0):
                    last_btn_status = 1  #������1�´ο��Լ�������
                if cap_num < class_num:
                    img = draw_string(img, 0, 200, "press boot key to cap "+class_names[cap_num], color=lcd.WHITE,scale=1, bg=lcd.RED)
                elif cap_num < class_num + sample_num:
                    img = draw_string(img, 0, 200, "boot key to cap sample{}".format(cap_num-class_num), color=lcd.WHITE,scale=1, bg=lcd.RED)
        # train and predict   ѵ����Ԥ��
        if train_status == 0:
            if cap_num >= class_num + sample_num:
                print("start train")
                img = draw_string(img, 30, 100, "training...", color=lcd.WHITE,scale=2, bg=lcd.RED)
                lcd.display(img)
                classifier.train()
                print("train end")
                train_status = 1
        else:
            res_index = -1
            try:
                res_index, min_dist = classifier.predict(img)  #��ģ�ͳ�Ԥ������
                print("{:.2f}".format(min_dist))  #����С��������λ
            except Exception as e:
                print("predict err:", e)
            if res_index >= 0 and min_dist < THRESHOLD : #������ֵ�����þ���һ�����ޣ�����С��11��ô�Ͳ�����ͨ����ģ��
                print("predict result:", class_names[res_index])
                uart_A.write('1\r\n')
                img = draw_string(img, 2, 2, class_names[res_index], color=lcd.WHITE,scale=2, bg=lcd.RED)
            else:
                uart_A.write('2\r\n')
                print("unknown, maybe:", class_names[res_index])
                img = draw_string(img, 2, 2, 'maybe {}'.format(class_names[res_index]), color=lcd.WHITE,scale=2, bg=lcd.RED)

        lcd.display(img)

    # You can save trained data to file system by:  ��ѵ�����ݱ��浽�ļ�ϵͳ
    classifier.save("3_classes.classifier")

################
def cheak_key():
    while(1):
        if KEY0.value()==0:
            indepent_learn()
        elif KEY1.value()==0:
            mask_learn()
        else :
            break

task_fd = kpu.load("/sd/fd.smodel")
task_ld = kpu.load("/sd/ld.smodel")
task_fe = kpu.load("/sd/fe.smodel")
clock = time.clock()

lcd.init()
sensor.reset()
sensor.set_pixformat(sensor.RGB565)
sensor.set_framesize(sensor.QVGA)
sensor.set_hmirror(0)
sensor.set_vflip(0)
sensor.run(1)
anchor = (1.889, 2.5245, 2.9465, 3.94056, 3.99987, 5.3658, 5.155437,
          6.92275, 6.718375, 9.01025)
dst_point = [(44, 59), (84, 59), (64, 82), (47, 105),
             (81, 105)]
a = kpu.init_yolo2(task_fd, 0.5, 0.3, 5, anchor)
img_lcd = image.Image()
img_face = image.Image(size=(128, 128))
a = img_face.pix_to_ai()
record_ftr = []
record_ftrs = []
names = ['Mr.1', 'Mr.2', 'Mr.3', 'Mr.4', 'Mr.5',
         'Mr.6', 'Mr.7', 'Mr.8', 'Mr.9', 'Mr.10']
ACCURACY = 85  #ģ����������ֵ
while (1):
    cheak_key()
    img = sensor.snapshot()
    clock.tick()
    code = kpu.run_yolo2(task_fd, img)
    if code:
        for i in code:
            a = img.draw_rectangle(i.rect())
            face_cut = img.cut(i.x(), i.y(), i.w(), i.h())
            face_cut_128 = face_cut.resize(128, 128) #�ҵ������������ó�128*128ͼ
            a = face_cut_128.pix_to_ai() #���ü���ͼ���ص�kpu����������
            fmap = kpu.forward(task_ld, face_cut_128) #�����Ѽ��ص�����ģ�͵�ָ������������Ŀ����������ͼ
            plist = fmap[:]
            le = (i.x() + int(plist[0] * i.w() - 10), i.y() + int(plist[1] * i.h()))  #����
            re = (i.x() + int(plist[2] * i.w()), i.y() + int(plist[3] * i.h()))  #����
            nose = (i.x() + int(plist[4] * i.w()), i.y() + int(plist[5] * i.h()))  #����
            lm = (i.x() + int(plist[6] * i.w()), i.y() + int(plist[7] * i.h())) # ������
            rm = (i.x() + int(plist[8] * i.w()), i.y() + int(plist[9] * i.h()))  #������
            a = img.draw_circle(le[0], le[1], 4)
            a = img.draw_circle(re[0], re[1], 4)
            a = img.draw_circle(nose[0], nose[1], 4)
            a = img.draw_circle(lm[0], lm[1], 4)
            a = img.draw_circle(rm[0], rm[1], 4)
            src_point = [le, re, nose, lm, rm]
            T = image.get_affine_transform(src_point, dst_point) #���������з����任��ת��Ϊ����ͼ��
            a = image.warp_affine_ai(img, img_face, T)
            a = img_face.ai_to_pix() #������ͼ������
            del (face_cut_128)  #ɾ����һ�����㴦����ͼ����Ϊ�˼���������һ��ͼ��
            fmap = kpu.forward(task_fe, img_face) #�������õ�����ͼ�����ؽ�������ģ����������������
            feature = kpu.face_encode(fmap[:])
            reg_flag = False
            scores = []
            for j in range(len(record_ftrs)):  #����forѭ�������þ��Ǳ���ÿ��������Ϣ�Ĳɼ�
                score = kpu.face_compare(record_ftrs[j], feature)
                scores.append(score)
            max_score = 0
            index = 0
            for k in range(len(scores)):   #��ģ���ܳ����ķ������и�ֵ��������ʾ��lcd��
                if max_score < scores[k]:
                    max_score = scores[k]
                    index = k
            if max_score > ACCURACY:
                uart_A.write('1\r\n')
                a = img.draw_string(i.x(), i.y(), ("%s :%2.1f" % (
                    names[index], max_score)), color=(0, 255, 0), scale=2)
            else:
                #uart_A.write('2\r\n')
                a = img.draw_string(i.x(), i.y(), ("X :%2.1f" % (
                    max_score)), color=(255, 0, 0), scale=2)
            if start_processing:
                record_ftr = feature  #���������洢������ÿ��������������Ϣ
                record_ftrs.append(record_ftr)  #������Ԫ�����ӽ�����ĩβ
                start_processing = False
            break
    a = lcd.display(img)
    gc.collect()

